# !/usr/bin/env python
# encoding: utf-8


"""
  @author: gaogao
  @file: Heater P-V.py
  @time: 2022/3/30 15:51
  @desc:
"""
import time
import datetime
from math import *
from ATS.application.LCR_App.U2722A_App import U2722A_App
from ATS.application.LCR_App.B2901BL_App import B2901BLApp
from ATS.application.LASER_App.L8164B_App import L8164B_App
from ATS.application.LASER_App.N7744C_App import N7744C_App
from ATS.datamanager import data_transfer_station as DataTransferStation
from utils.read_ini_2_dict import get_option_value, get_refrence_data, get_loss_data, get_pd_loss_sys
from utils.get_wave_power_util import get_wave_power_value, get_wave_loss_gc
from ATS.util.calcu_step_util import get_points_list, get_step_list, get_step_array
from ATS.util.save_data_2_csv_util import saveData2CSV
from utils.find3db_util import find3db_util
from collections import namedtuple
from matplotlib import pyplot as plt
import numpy as np
from utils.log_util import loggings
from gui.driver.var_manager import VarManager
from utils.save_setting_2_ini import save_ini_info
from gui import constant
from utils.save_data_2_file import save_data_2_file_exper, save_spot_data_2_file
from utils.polynomial_fit_util import normalized, polynomial_fit
from database.db_interface import insert_chip

try:
    from ATS.switchdriver.LightSwitchDriver import LightSwitchDriver
except Exception as e:
    raise ValueError("光开关不在线，请重启光开关！")


# 通道1设置扫描电压，通道2、3固定电压
def er2_test(para):
    try:
        time_stamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        delay_time = 2
        voltage_vpp_list = get_step_array(0.0, 1.5, 0.05)
        result_data = namedtuple("result_data",
                                 "file_name,curve_type,y_invert,title,x_label,y_label,test_result")
        total_time = 10
        inst_light = LightSwitchDriver.instance()
        if not inst_light.is_connected:
            inst_light.connect()
        inst_light.set_all_switch("01", "05")
        time.sleep(1)
        loggings.info("switch" + str(inst_light.query_switch()))
        

        l8164 = para.l8164_setting
        n7744 = para.n7744_setting
        l8164b_app = L8164B_App(l8164["title"])
        n7744c_app = N7744C_App(n7744["title"])

        l8164b_app.set_power_unit()
        l8164b_app.set_power_value(float(l8164["power_value"]))
        wave_length = l8164["wave_length"][0]
        l8164b_app.set_wave_length(int(wave_length))
        time.sleep(4)
        p_in = l8164b_app.fetch_input_power()
        l8164b_app.output_on()
        n7744c_app.reset()
        time.sleep(2)
        n7744c_app.set_power_wave_length()
        n7744c_app.set_power_unit()
        n7744c_app.set_power_range_auto()
        n7744c_app.set_power_atim(n7744["atim"])

        time.sleep(2)
        smu_setting = para.smu_setting
        smu_dict = {}
        # smu voltage_list channel nplc current_range current_limit //mode sweep
        spci_meas_list = []
        mes_plans_dict = {"var_func": None, "var1": None, "measure_func": None}
        init_set_list = []
        b29smu_dict = {}
        init_b29_set_list = []
        b29_set_voltage_list = []
        b29_get_current_list = []
        b29_measure_list = []
        if smu_setting:
            for smu in smu_setting:
                # if smu["smu"] == "SMU10":
                if smu["smu"] == "SMU10" or smu["smu"] == "SMU11":
                    if smu["smu"] not in b29smu_dict:
                        b29smu_dict[smu["smu"]] = B2901BLApp(smu["smu"])
                    current_inst = b29smu_dict[smu["smu"]]
                    bias_voltage = smu["voltage_list"][0]
                    current_limit = smu["current_limit"]
                    init_b29_set_list += [lambda inst_obj=current_inst, voltage=bias_voltage,
                                                 current_limit=current_limit: inst_obj.spot_setting(
                        float(voltage), current_limit)]

                    b29_set_voltage_list += [
                        lambda x, inst_obj=current_inst, voltage=bias_voltage: inst_obj.set_voltage(
                            float(voltage) + float(x))]

                    b29_get_current_list += [lambda inst_obj=current_inst: inst_obj.get_current()]
                    b29_measure_list += [lambda inst_obj=current_inst: inst_obj.trigger_measure()]

                else:
                    if smu["smu"] not in smu_dict:
                        smu_dict[smu["smu"]] = U2722A_App(smu["smu"])
                    current_inst = smu_dict[smu["smu"]]
                    current_voltage_list = smu["voltage_list"]
                    current_channel = smu["channel"]
                    current_channel_nplc = smu["nplc"]
                    current_channel_range = smu["current_range"]
                    current_channel_limit = smu["current_limit"]
                    init_set_list += [lambda y=current_channel,
                                             inst_obj=current_inst: inst_obj.output_state(0, y),
                                      lambda x=current_channel, inst_obj=current_inst: inst_obj.set_voltage_range(x),
                                      lambda x=current_channel_nplc, y=current_channel,
                                             inst_obj=current_inst: inst_obj.set_current_nplc(x, y),
                                      lambda x=current_channel_range, y=current_channel_limit,
                                             z=current_channel, inst_obj=current_inst: inst_obj.set_current_range_limit(
                                          x,
                                          y,
                                          z)]
                    if smu["mode"] == "sweep":
                        mes_plans_dict["var1"] = current_voltage_list
                        mes_plans_dict["var_func"] = lambda volt, x=current_channel, inst_obj=current_inst: \
                            inst_obj.set_voltage(volt, x)
                        mes_plans_dict["measure_func"] = lambda x=current_channel, inst_obj=current_inst: abs(
                            float(inst_obj.get_current_value(x)))
                    else:
                        init_set_list += [lambda x=current_voltage_list[0], y=current_channel,
                                                 inst_obj=current_inst: inst_obj.set_voltage(x, y)]

                        spci_meas_list += [lambda x, y=current_channel, z=current_voltage_list[0],
                                                  inst_obj=current_inst: inst_obj.set_voltage(
                            float(x) + float(z), y)]

                    if smu["is_measure"]:
                        mes_plans_dict["measure_func"] = lambda x=current_channel, inst_obj=current_inst: abs(
                            float(inst_obj.get_current_value(x)))

                    init_set_list += [lambda y=current_channel,
                                             inst_obj=current_inst: inst_obj.output_state(1, y)]

            # set init
            for inst_set_func in init_set_list:
                inst_set_func()  # 设置SMU
            for b29_set_func in init_b29_set_list:
                b29_set_func()
            # read mes
            power_list = []
            mes_func = mes_plans_dict["measure_func"]
            if mes_func:  # mes func不为空
                if mes_plans_dict["var1"]:  # sweep
                    for volt in mes_plans_dict["var1"]:
                        mes_plans_dict["var_func"](volt)
                        for b29_measre in b29_measure_list:
                            b29_measre()
                        time.sleep(0.1)
                        power_value = n7744c_app.get_power_value()
                        power_list.append(float(power_value))
            sweep_list = [float(i) for i in mes_plans_dict["var1"]]

            wave_power_tuple = get_refrence_data()
            loss_section = f"{VarManager().diex}-{VarManager().diey}-{VarManager().subdie_name.split('_')[0]}_GC"
            loss_gc = get_wave_loss_gc(l8164["wave_length"][0], get_loss_data(section=loss_section))
            loss_sys = get_wave_power_value(l8164["wave_length"][0], wave_power_tuple)
            pd_loss_sys = get_wave_power_value(l8164["wave_length"][0], get_pd_loss_sys())
            title_section = f"{VarManager().diex}-{VarManager().diey}&{VarManager().subdie_name}&{para.subdie}&{bias_voltage}V_{wave_length}"
            title_path = f"{VarManager().save_folder}/{VarManager().wafer_name}"
            # save_spot_data_2_file(f"{title_path}.csv", f"{title_section}_vpi_loss_gc", str(loss_gc))
            for num_index in range(len(power_list)):
                power_list[num_index] = round(
                    (float(p_in) - power_list[num_index] - float(loss_sys) - 2 * float(loss_gc)), 3)

            vpi_pv = f"{title_section}_PVCurve_{time_stamp}"

            curve_title = f"{title_section}_PVCurve_{time_stamp}"
            x_label = "Voltage"
            y_label = "Power"

            DataTransferStation.get_DTS().postarraydata("test", result_data(vpi_pv, 'line', True,
                                                                            curve_title, x_label, y_label,
                                                                            (sweep_list, power_list)))

            chip_name = f"{VarManager().subdie_name}"
            wafer_name = VarManager().wafer_name
            shot_position = f"{VarManager().diex},{VarManager().diey}"
            shot_number = VarManager().shot_number
            chip_number = VarManager().subdie_number
            polarization_type = VarManager().polarization_type
            # measurement_parameter = f"{para.subdie}_{para.test_item}_PVCurve"
            measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_PVCurve"
            measurement_data = str([sweep_list, power_list])
            data_unit = ""
            data_type = "Curve"
            insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                        measurement_parameter, measurement_data, data_unit, data_type)

            datas = find3db_util(sweep_list, power_list)
            if datas and len(datas) == 4:
                # 设置bias电压
                if datas[0] != "" and datas[0] is not None:
                    half_slope_voltage = datas[0]
                    if mes_plans_dict["var1"]:
                        mes_plans_dict["var_func"](half_slope_voltage)
                    voltage_vpp_list = [float(i) for i in voltage_vpp_list]
                    save_spot_data_2_file(f"{title_path}.csv", f"{title_section}_Heater_3dB", str(half_slope_voltage))
                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_Heater_3dB"
                    measurement_data = str(half_slope_voltage)
                    data_unit = "V"
                    data_type = "Spot"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)

                    p1_list, b2901_current_list, b2902_current_list = v_pi_test3(b29_set_voltage_list, bias_voltage,
                                                                                 delay_time, n7744c_app, total_time,
                                                                                 voltage_vpp_list, b29_get_current_list,b29_measure_list)

                    time.sleep(6)
                    p2_list, b2901_current_list, b2902_current_list = v_pi_test4(b29_set_voltage_list, bias_voltage,
                                                                                 delay_time, n7744c_app, total_time,
                                                                                 voltage_vpp_list, b29_get_current_list,b29_measure_list)

                    time.sleep(2)
                    # p2_list = v_pi_test2(b2901, bias_voltage, delay_time, n7744c_app, spci_meas_list, total_time,
                    #                      voltage_vpp_list)
                    #
                    # p1_list = v_pi_test(b2901, bias_voltage, delay_time, n7744c_app, spci_meas_list, total_time,
                    #                     voltage_vpp_list)

                    save_spot_data_2_file(f"{title_path}.csv", f"{title_section}_IL", str(datas[3]))

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_IL"
                    measurement_data = str(datas[3])
                    data_unit = "dB"
                    data_type = "Spot"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)

                    p0 = float(datas[3]) - float(datas[1])

                    save_spot_data_2_file(f"{title_path}.csv", f"{title_section}_Power_3dB",
                                          str(p0))

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_Power_3dB"
                    measurement_data = str(p0)
                    data_unit = "dB"
                    data_type = "Spot"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)

                    p11_list = normalized(p0, p1_list)
                    p22_list = normalized(p0, p2_list)

                    p11_result = just_flag(p11_list)

                    p11_file_name = f"{title_section}_power_left_arm_{p11_result}_{time_stamp}"

                    curve_title = f"{title_section}_power_left_arm_{p11_result}_{time_stamp}"
                    x_label = "Voltage"
                    y_label = "Power"
                    post_array_data(curve_title, p11_list, result_data, p11_file_name, voltage_vpp_list, x_label,
                                    y_label)

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_power_left_arm_{p11_result}"
                    measurement_data = str([voltage_vpp_list, p11_list])
                    data_unit = ""
                    data_type = "Curve"

                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)

                    time.sleep(2)
                    p22_result = just_flag(p22_list)
                    p11_file_name = f"{title_section}_power_right_arm_{p22_result}_{time_stamp}"

                    curve_title = f"{title_section}_power_right_arm_{p22_result}_{time_stamp}"
                    x_label = "Voltage"
                    y_label = "Power"

                    post_array_data(curve_title, p22_list, result_data, p11_file_name, voltage_vpp_list, x_label,
                                    y_label)

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_power_right_arm_{p22_result}"
                    measurement_data = str([voltage_vpp_list, p22_list])
                    data_unit = ""
                    data_type = "Curve"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)

                    voltage_vpp_list = [float(i) for i in voltage_vpp_list]

                    v_pi2, z2 = polynomial_fit(voltage_vpp_list, p22_list)
                    v_pi1, z1 = polynomial_fit(voltage_vpp_list, p11_list)

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_Vpi_left_arm_abc"
                    measurement_data = str(z1)
                    data_unit = ""
                    data_type = "Spot"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_Vpi_right_arm_abc"
                    measurement_data = str(z2)
                    data_unit = ""
                    data_type = "Spot"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_Vpi_left_arm_{p11_result}"
                    measurement_data = str(v_pi1)
                    data_unit = "V"
                    data_type = "Spot"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)
                    save_spot_data_2_file(f"{title_path}.csv", f"{title_section}_left_arm_{p11_result}",
                                          str(v_pi1))

                    measurement_parameter = f"{para.subdie}&{bias_voltage}V_{wave_length}_Vpi_right_arm_{p22_result}"
                    measurement_data = str(v_pi2)
                    data_unit = "V"
                    data_type = "Spot"
                    insert_chip(chip_name, wafer_name, shot_position, shot_number, chip_number, polarization_type,
                                measurement_parameter, measurement_data, data_unit, data_type)
                    save_spot_data_2_file(f"{title_path}.csv", f"{title_section}_right_arm_{p22_result}",
                                          str(v_pi2))
            else:
                loggings.info("No Value")
    except Exception as e:
        loggings.error("er2_test" + str(e))
    finally:
        l8164b_app.set_wave_length(1330)
        close_all_channel()


def post_array_data(curve_title, p22_list, result_data, title_section, voltage_vpp_list, x_label, y_label):
    DataTransferStation.get_DTS().postarraydata("test",
                                                result_data(title_section, 'line', False,
                                                            curve_title, x_label, y_label,
                                                            (voltage_vpp_list, p22_list)))


def just_flag(vp_list):
    if float(vp_list[-1]) > float(vp_list[0]):
        flag = "up"
    else:
        flag = "down"
    return flag


def v_pi_test(b2901, bias_voltage, delay_time, n7744c_app, spci_meas_list, total_time, voltage_vpp_list):
    p_total_list = []
    for voltage in voltage_vpp_list:
        spci_meas_list[-1](0)
        b2901.set_voltage(float(voltage) + float(bias_voltage))
        p_list = []
        for _ in range(total_time):
            p = n7744c_app.get_power_value()
            p_list.append(float(p))
            time.sleep(0.05)
        #
        spci_meas_list[-1](-float(bias_voltage))
        b2901.set_voltage(0)
        min_p = min(p_list)
        max_p = max(p_list)
        p_list.remove(min_p)
        p_list.remove(max_p)
        p = sum(p_list) / len(p_list)
        p_total_list.append(p)
        time.sleep(delay_time)
    return p_total_list


def v_pi_test3(b29_set_voltage_list, bias_voltage, delay_time, n7744c_app, total_time, voltage_vpp_list,
               b29_get_current_list,b29_measure_list):
    p_total_list = []
    b2901_current = []
    b2902_current = []
    for voltage in voltage_vpp_list:
        b29_set_voltage_list[0](float(voltage))
        b29_set_voltage_list[1](0)
        for b29_measre in b29_measure_list:
            b29_measre()
        time.sleep(0.05)
        b2901_current.append(float(b29_get_current_list[0]()))
        b2902_current.append(float(b29_get_current_list[1]()))
        p_list = []
        for _ in range(total_time):
            p = n7744c_app.get_power_value()
            p_list.append(float(p))
            time.sleep(0.05)
        min_p = min(p_list)
        max_p = max(p_list)
        p_list.remove(min_p)
        p_list.remove(max_p)
        p = sum(p_list) / len(p_list)
        p_total_list.append(p)
        b29_set_voltage_list[0](-float(bias_voltage))
        b29_set_voltage_list[1](-float(bias_voltage))
        time.sleep(delay_time)
    return p_total_list, b2902_current, b2901_current


def v_pi_test4(b29_set_voltage_list, bias_voltage, delay_time, n7744c_app, total_time, voltage_vpp_list,
               b29_get_current_list,b29_measure_list):
    p_total_list = []
    b2901_current = []
    b2902_current = []
    for voltage in voltage_vpp_list:
        b29_set_voltage_list[1](float(voltage))
        b29_set_voltage_list[0](0)
        for b29_measre in b29_measure_list:
            b29_measre()
        time.sleep(0.05)
        b2901_current.append(float(b29_get_current_list[0]()))
        b2902_current.append(float(b29_get_current_list[1]()))
        p_list = []
        for _ in range(total_time):
            p = n7744c_app.get_power_value()
            p_list.append(float(p))
            time.sleep(0.05)
        min_p = min(p_list)
        max_p = max(p_list)
        p_list.remove(min_p)
        p_list.remove(max_p)
        p = sum(p_list) / len(p_list)
        p_total_list.append(p)
        b29_set_voltage_list[1](-float(bias_voltage))
        b29_set_voltage_list[0](-float(bias_voltage))
        time.sleep(delay_time)
    return p_total_list, b2901_current, b2902_current


def v_pi_test2(b2901, bias_voltage, delay_time, n7744c_app, spci_meas_list, total_time, voltage_vpp_list):
    p_total_list = []
    for voltage in voltage_vpp_list:
        spci_meas_list[-1](float(voltage))
        b2901.set_voltage(bias_voltage)
        p_list = []
        for _ in range(total_time):
            p = n7744c_app.get_power_value()
            p_list.append(float(p))
            time.sleep(0.1)
        #

        spci_meas_list[-1](-float(bias_voltage))
        b2901.set_voltage(0)
        min_p = min(p_list)
        max_p = max(p_list)
        p_list.remove(min_p)
        p_list.remove(max_p)

        p = sum(p_list) / len(p_list)
        p_total_list.append(p)
        time.sleep(delay_time)
    return p_total_list


def close_all_channel():
    U2722A_App("SMU2").close_all_channel()
    U2722A_App("SMU3").close_all_channel()
    U2722A_App("SMU4").close_all_channel()
    B2901BLApp("SMU10").output_close()
    B2901BLApp("SMU11").output_close()


er2_test(locals()['para'])
